import time
import numpy as np
from rtde_control import RTDEControlInterface
from rtde_receive import RTDEReceiveInterface

from ezgl import Matrix4x4
from .robotiq import RobotiqGripper


class UR5:
    def __init__(self, arm_ip, with_gripper=True):
        arm_ip = arm_ip
        self.rtde_c = RTDEControlInterface(arm_ip, frequency=500)
        self.rtde_r = RTDEReceiveInterface(arm_ip)
        self.with_gripper = with_gripper

        if with_gripper:
            self.gripper = RobotiqGripper()
            self.gripper.connect(arm_ip, port=63352, socket_timeout=0.5)

    def moveL(self, ee_pose: Matrix4x4, vel=0.25, acc=0.5, asynchronous=False):
        self.rtde_c.moveL(ee_pose.toTransRotVector(), vel, acc, asynchronous)

    def servoL(self, ee_pose: Matrix4x4, vel=0.25, acc=0.5, period=0.008, lookahead=0.2, gain=100):
        self.rtde_c.servoL(ee_pose.toTransRotVector(), vel, acc, period, lookahead, gain)

    def moveJ(self, joints_rad, vel=0.25, acc=0.5, asynchronous=False):
        self.rtde_c.moveJ(joints_rad, vel, acc, asynchronous)

    def servoJ(self, joints_rad, vel=0.25, acc=0.5, period=0.008, lookahead=0.2, gain=100):
        self.rtde_c.servoJ(joints_rad, vel, acc, period, lookahead, gain)

    def get_pose(self):
        pose = self.rtde_r.getActualTCPPose() # xyz + orientation in rotation vector format
        return Matrix4x4.fromTransRotVector(*pose)
    
    def get_joints_rad(self):
        angles = self.rtde_r.getActualQ()
        return angles
    
    def stop(self, control_mode="joint"):
        if control_mode == "joint":
            self.rtde_c.stopJ()
        elif control_mode == "ee":
            self.rtde_c.stopL()
        else:
            raise ValueError("Invalid control mode")

    def gripper_move(self, gripper_distance, speed=125, force=255):
        command = int((1 - gripper_distance / 128) * 255 * 0.85)
        self.gripper.move(command, speed=speed, force=force)

    def close(self):
        if self.with_gripper:
            self.gripper.disconnect()
        self.rtde_c.disconnect()
        self.rtde_r.disconnect()

    def teach_mode(self):
        self.rtde_c.teachMode()

    def end_teach_mode(self):
        self.rtde_c.endTeachMode()



if __name__=="__main__":
    import yaml
    config_path = r"./configs/config.yaml"
    with open(config_path, "r") as f:
        config = yaml.safe_load(f)
    arm = UR5(config)
    while True:
        try:
            pose = arm.xarm_move([0.2,0.35,0.6,-np.pi/2,np.pi/2,0])
            # pose = arm.xarm_move([0.2,0.35,0.6,4.712,0,0])
            arm_pose = arm.get_pose()
            xyz = arm_pose[:3]
            ori = np.degrees(arm_pose[3:])
            print("Current ori:", ori)
            arm.gripper_move(0)
            # print(arm.get_angles())
            time.sleep(1)
        except KeyboardInterrupt:
            print("Stopping arm...")
            arm.close()
            break